from __future__ import annotations


# imports
from matplotlib.pyplot import _get_pyplot_commands

# Dummy colormaps for test, since matplotlib.cm._colormaps is not available
_colormaps = {"viridis", "plasma", "inferno", "magma"}

colormaps = _colormaps

# --- Unit tests for _get_pyplot_commands ---


# Helper function to add/remove functions to the module globals for testing
def add_function_to_globals(name, func):
    globals()[name] = func


def remove_function_from_globals(name):
    if name in globals():
        del globals()[name]


# Basic test cases


def test_basic_single_function():
    """Test with a single public function in globals."""

    def foo():
        pass

    add_function_to_globals("foo", foo)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 118μs -> 98.7μs (19.9% faster)
    remove_function_from_globals("foo")


def test_basic_multiple_functions():
    """Test with multiple public functions in globals."""

    def bar():
        pass

    def baz():
        pass

    add_function_to_globals("bar", bar)
    add_function_to_globals("baz", baz)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 117μs -> 97.4μs (21.0% faster)
    remove_function_from_globals("bar")
    remove_function_from_globals("baz")


def test_exclude_private_functions():
    """Test that functions starting with '_' are excluded."""

    def _private():
        pass

    add_function_to_globals("_private", _private)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 116μs -> 96.4μs (21.3% faster)
    remove_function_from_globals("_private")


def test_exclude_colormap_functions():
    """Test that colormap-named functions are excluded."""
    for cmap in colormaps:

        def cmap_func():
            pass

        add_function_to_globals(cmap, cmap_func)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 116μs -> 96.3μs (21.4% faster)
    for cmap in colormaps:
        remove_function_from_globals(cmap)


def test_exclude_get_plot_commands_and_colors():
    """Test that 'get_plot_commands' and 'colors' are excluded."""

    def get_plot_commands():
        pass

    def colors():
        pass

    add_function_to_globals("get_plot_commands", get_plot_commands)
    add_function_to_globals("colors", colors)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 116μs -> 96.2μs (21.3% faster)
    remove_function_from_globals("get_plot_commands")
    remove_function_from_globals("colors")


def test_exclude_non_function_globals():
    """Test that non-function globals are excluded."""
    globals()["not_a_func"] = 123
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 117μs -> 96.2μs (21.6% faster)
    del globals()["not_a_func"]


def test_function_with_weird_names():
    """Test that functions with weird names are handled correctly."""

    def func123():
        pass

    def func_with_underscore():
        pass

    def funcWithCaps():
        pass

    add_function_to_globals("func123", func123)
    add_function_to_globals("func_with_underscore", func_with_underscore)
    add_function_to_globals("funcWithCaps", funcWithCaps)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 117μs -> 97.1μs (21.3% faster)
    remove_function_from_globals("func123")
    remove_function_from_globals("func_with_underscore")
    remove_function_from_globals("funcWithCaps")


def test_function_from_other_module_excluded():
    """Test that a function from another module is excluded."""

    def external_func():
        pass

    # Simulate external function by changing its __module__ attribute
    external_func.__module__ = "external"
    add_function_to_globals("external_func", external_func)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 117μs -> 96.5μs (21.6% faster)
    remove_function_from_globals("external_func")


def test_function_with_leading_and_trailing_spaces():
    """Test that function names with spaces are handled (shouldn't happen, but test anyway)."""

    def weird_func():
        pass

    # Add with a name containing spaces
    add_function_to_globals(" weird_func ", weird_func)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 116μs -> 95.6μs (21.5% faster)
    remove_function_from_globals(" weird_func ")


def test_function_named_excluded_but_private():
    """Test that a function named like an exclusion but private is excluded."""

    def _colors():
        pass

    add_function_to_globals("_colors", _colors)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 115μs -> 96.7μs (19.4% faster)
    remove_function_from_globals("_colors")


def test_function_with_non_str_name():
    """Test that globals with non-str keys are ignored."""

    def dummy():
        pass

    globals()[42] = dummy
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 118μs -> 97.0μs (21.8% faster)
    del globals()[42]


# Large scale test cases


def test_many_functions():
    """Test with a large number of functions (performance and correctness)."""
    func_names = [f"func_{i}" for i in range(500)]
    for name in func_names:

        def f():
            pass

        add_function_to_globals(name, f)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 118μs -> 97.4μs (21.3% faster)
    for name in func_names:
        pass
    for name in func_names:
        remove_function_from_globals(name)


def test_many_excluded_functions():
    """Test with a large number of excluded functions (colormaps)."""
    excluded_names = [f"colormap_{i}" for i in range(500)]
    for name in excluded_names:

        def f():
            pass

        add_function_to_globals(name, f)
    # Add to colormaps for exclusion
    colormaps.update(excluded_names)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 116μs -> 95.9μs (21.1% faster)
    for name in excluded_names:
        pass
    for name in excluded_names:
        remove_function_from_globals(name)
    # Remove from colormaps
    for name in excluded_names:
        colormaps.discard(name)


def test_mixture_of_functions_and_non_functions():
    """Test with a mixture of functions and non-functions in globals."""
    func_names = [f"mix_func_{i}" for i in range(100)]
    non_func_names = [f"mix_non_func_{i}" for i in range(100)]
    for name in func_names:

        def f():
            pass

        add_function_to_globals(name, f)
    for name in non_func_names:
        globals()[name] = name
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 116μs -> 95.9μs (21.9% faster)
    for name in func_names:
        pass
    for name in non_func_names:
        pass
    for name in func_names:
        remove_function_from_globals(name)
    for name in non_func_names:
        del globals()[name]


def test_performance_large_scale():
    """Test that function works efficiently with near 1000 functions."""
    func_names = [f"perf_func_{i}" for i in range(950)]
    for name in func_names:

        def f():
            pass

        add_function_to_globals(name, f)
    codeflash_output = _get_pyplot_commands()
    result = codeflash_output  # 117μs -> 97.1μs (21.2% faster)
    for name in func_names:
        remove_function_from_globals(name)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.